#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <stdbool.h>
#include <unistd.h>
#include <time.h>
#include <sys/mman.h>
#include <security/pam_appl.h>
#include "raylib.h"

// ================================================================================
// Configuration Macros - Modify these to adjust behavior
// ================================================================================
#define ENABLE_DEBUG_LOG 0 // 1 = Enable debug output to stderr, 0 = Disable
#define ALLOW_ESC_QUIT 0   // 1 = ESC quits (debug), 0 = ESC clears input (system)
#define WINDOW_TITLE "System Screen Locker"
#define FONT_SIZE_MAIN 40
#define FONT_SIZE_SUB 30
#define FONT_SIZE_INDICATOR 28
#define MAX_PASSWORD_LEN 512
#define MAX_FAILED_ATTEMPTS 5
#define PADDING 15
#define SYSTEM_FONT_PATH "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf"

// Color definitions
#define COLOR_BACKGROUND (Color){0, 0, 0, 255}  // Black
#define COLOR_TEXT (Color){255, 255, 255, 255}  // White
#define COLOR_ERROR (Color){255, 0, 0, 255}     // Red
#define COLOR_SUCCESS (Color){0, 255, 0, 255}   // Green
#define COLOR_INPUT_BG (Color){40, 40, 40, 255} // Dark gray
#define COLOR_PROMPT (Color){255, 204, 0, 255}  // Amber

// ================================================================================
// Global State
// ================================================================================
char password[MAX_PASSWORD_LEN] = {0};
int input_pos = 0;
int failed_attempts = 0;
bool is_locked = true;
bool show_help = false;
size_t cursor = 0;
Font app_font;
Font default_font;
bool is_custom_font = false;

// ================================================================================
// PAM Authentication Callbacks
// ================================================================================
static int pam_conv_callback(int num_msg, const struct pam_message **msg,
                             struct pam_response **resp, void *appdata_ptr)
{
    (void)appdata_ptr;
    *resp = calloc(num_msg, sizeof(struct pam_response));
    if (!*resp)
        return PAM_BUF_ERR;

    for (int i = 0; i < num_msg; i++)
    {
        if (msg[i]->msg_style == PAM_PROMPT_ECHO_OFF)
        {
            resp[i]->resp = strdup(password);
            resp[i]->resp_retcode = 0;
        }
    }
    return PAM_SUCCESS;
}

static bool verify_password()
{
    if (input_pos == 0)
        return false;

    struct pam_conv conv = {pam_conv_callback, NULL};
    pam_handle_t *pam_handle = NULL;
    const char *username = getlogin();
    int ret = pam_start("login", username, &conv, &pam_handle);

    if (ret != PAM_SUCCESS)
    {
#if ENABLE_DEBUG_LOG
        fprintf(stderr, "PAM Init Fail: %s\n", pam_strerror(pam_handle, ret));
#endif
        pam_end(pam_handle, ret);
        return false;
    }

    ret = pam_authenticate(pam_handle, 0);
    pam_end(pam_handle, ret);

    if (ret == PAM_SUCCESS)
    {
#if ENABLE_DEBUG_LOG
        printf("Auth Success!\n");
#endif
        return true;
    }
    else
    {
#if ENABLE_DEBUG_LOG
        fprintf(stderr, "Auth Fail: %s\n", pam_strerror(NULL, ret));
#endif
        failed_attempts++;
        return false;
    }
}

// ================================================================================
// Password Management
// ================================================================================
static void clear_password()
{
    memset(password, 0, sizeof(password));
    input_pos = 0;
    cursor = 0;
}

// ================================================================================
// Input Handling
// ================================================================================
static void handle_input()
{
    if (!is_locked)
        return;

    bool ctrl_pressed = IsKeyDown(KEY_LEFT_CONTROL) || IsKeyDown(KEY_RIGHT_CONTROL);

    // F1: Toggle help
    if (IsKeyPressed(KEY_F1))
    {
        show_help = !show_help;
    }

    // ESC: Quit (debug) or clear input (system)
    if (IsKeyPressed(KEY_ESCAPE))
    {
#if ALLOW_ESC_QUIT
        is_locked = false;
#else
        clear_password();
#endif
        return;
    }

    // Backspace: Delete previous character
    if (IsKeyPressed(KEY_BACKSPACE) && cursor > 0)
    {
        cursor--;
        memmove(&password[cursor], &password[cursor + 1], MAX_PASSWORD_LEN - cursor - 1);
        input_pos--;
        return;
    }

    // Enter: Submit password
    if (IsKeyPressed(KEY_ENTER) || IsKeyPressed(KEY_KP_ENTER))
    {
        if (verify_password())
        {
            is_locked = false;
        }
        else
        {
            clear_password();
        }
        return;
    }

    // Character input
    int key = GetCharPressed();
    while (key > 0)
    {
        if (key >= 32 && key <= 126 && input_pos < MAX_PASSWORD_LEN - 1)
        {
            memmove(&password[cursor + 1], &password[cursor], MAX_PASSWORD_LEN - cursor - 1);
            password[cursor++] = (char)key;
            input_pos++;
#if ENABLE_DEBUG_LOG
            fprintf(stderr, "Input: %c (Length: %d)\n", (char)key, input_pos);
#endif
        }
        key = GetCharPressed();
    }

    // Cursor navigation (arrow keys + Ctrl shortcuts)
    if (IsKeyPressed(KEY_RIGHT) || (ctrl_pressed && IsKeyPressed(KEY_F)))
        cursor = (cursor < (size_t)input_pos) ? cursor + 1 : cursor;

    if (IsKeyPressed(KEY_LEFT) || (ctrl_pressed && IsKeyPressed(KEY_B)))
        cursor = (cursor > 0) ? cursor - 1 : cursor;

    if (IsKeyPressed(KEY_HOME) || (ctrl_pressed && IsKeyPressed(KEY_A)))
        cursor = 0;

    if (IsKeyPressed(KEY_END) || (ctrl_pressed && IsKeyPressed(KEY_E)))
        cursor = input_pos;

    // Delete forward (Delete key + Ctrl+D)
    if ((IsKeyPressed(KEY_DELETE) || (ctrl_pressed && IsKeyPressed(KEY_D))) && cursor < (size_t)input_pos)
    {
        memmove(&password[cursor], &password[cursor + 1], MAX_PASSWORD_LEN - cursor - 1);
        input_pos--;
    }

    // Ctrl+K: Clear from cursor to end
    if (ctrl_pressed && IsKeyPressed(KEY_K))
    {
        password[cursor] = '\0';
        input_pos = cursor;
    }
}

// ================================================================================
// UI Rendering
// ================================================================================
static void draw_ui()
{
    BeginDrawing();
    ClearBackground(COLOR_BACKGROUND);

    int screen_w = GetScreenWidth();
    int screen_h = GetScreenHeight();
    Vector2 center = {screen_w / 2.0f, screen_h / 2.0f};

    // Title
    const char *title = "System Screen Locked";
    Vector2 title_size = MeasureTextEx(app_font, title, FONT_SIZE_MAIN, 1);
    DrawTextEx(app_font, title,
               (Vector2){center.x - title_size.x / 2, center.y - 120},
               FONT_SIZE_MAIN, 1, COLOR_TEXT);

    // Password input box
    int input_box_y = center.y - 40;
    DrawRectangle(PADDING, input_box_y, screen_w - 2 * PADDING,
                  FONT_SIZE_SUB + 12, COLOR_INPUT_BG);

    // Prompt text ("Password: ")
    const char *prompt = "Password: ";
    Vector2 prompt_size = MeasureTextEx(app_font, prompt, FONT_SIZE_SUB, 1);
    DrawTextEx(app_font, prompt,
               (Vector2){PADDING + 8, input_box_y + 6},
               FONT_SIZE_SUB, 1, COLOR_PROMPT);

    // Password indicator (*)
    char pass_indicator[MAX_PASSWORD_LEN + 1] = {0};
    memset(pass_indicator, '*', input_pos);
    DrawTextEx(app_font, pass_indicator,
               (Vector2){PADDING + 8 + prompt_size.x, input_box_y + 6},
               FONT_SIZE_SUB, 1, COLOR_SUCCESS);

    // Blinking cursor
    static double last_cursor_time = 0;
    static bool cursor_visible = true;
    double current_time = GetTime();
    if (current_time - last_cursor_time > 0.5)
    {
        cursor_visible = !cursor_visible;
        last_cursor_time = current_time;
    }
    if (cursor_visible)
    {
        char cursor_substr[MAX_PASSWORD_LEN];
        strncpy(cursor_substr, pass_indicator, cursor);
        cursor_substr[cursor] = '\0';
        float cursor_x = PADDING + 8 + prompt_size.x +
                         MeasureTextEx(app_font, cursor_substr, FONT_SIZE_SUB, 1).x;
        DrawRectangle(cursor_x, input_box_y + 6, 2, FONT_SIZE_SUB, COLOR_TEXT);
    }

    // Separator line
    int separator_y = input_box_y + FONT_SIZE_SUB + 20;
    DrawLine(PADDING, separator_y, screen_w - PADDING, separator_y,
             (Color){80, 80, 80, 255});

    // Failed attempts indicator
    if (failed_attempts > 0)
    {
        char fail_text[64];
        snprintf(fail_text, sizeof(fail_text), "Failed: %d/%d",
                 failed_attempts, MAX_FAILED_ATTEMPTS);
        Vector2 fail_size = MeasureTextEx(app_font, fail_text, FONT_SIZE_SUB - 2, 1);
        DrawTextEx(app_font, fail_text,
                   (Vector2){center.x - fail_size.x / 2, separator_y + 20},
                   FONT_SIZE_SUB - 2, 1, COLOR_ERROR);

        if (failed_attempts >= MAX_FAILED_ATTEMPTS)
        {
            const char *block_text = "Temporarily Blocked";
            Vector2 block_size = MeasureTextEx(app_font, block_text, FONT_SIZE_SUB - 2, 1);
            DrawTextEx(app_font, block_text,
                       (Vector2){center.x - block_size.x / 2, separator_y + 50},
                       FONT_SIZE_SUB - 2, 1, COLOR_ERROR);
        }
    }

    // Help text (F1 toggle)
    if (show_help)
    {
        const char *esc_behavior = ALLOW_ESC_QUIT ? "Quit (Debug)" : "Clear Input";
        char help_line1[128];
        snprintf(help_line1, sizeof(help_line1), "ESC: %s | Enter: Submit", esc_behavior);
        const char *help_line2 = "Ctrl+A: Home | Ctrl+E: End | Ctrl+F/B: Left/Right";
        const char *help_line3 = "Backspace: Delete | Ctrl+D: Delete Forward | Ctrl+K: Clear Rest";

        int help_y = screen_h - 100;
        Vector2 h1_size = MeasureTextEx(app_font, help_line1, FONT_SIZE_SUB - 4, 1);
        DrawTextEx(app_font, help_line1,
                   (Vector2){center.x - h1_size.x / 2, help_y},
                   FONT_SIZE_SUB - 4, 1, COLOR_TEXT);

        Vector2 h2_size = MeasureTextEx(app_font, help_line2, FONT_SIZE_SUB - 4, 1);
        DrawTextEx(app_font, help_line2,
                   (Vector2){center.x - h2_size.x / 2, help_y + 30},
                   FONT_SIZE_SUB - 4, 1, COLOR_TEXT);

        Vector2 h3_size = MeasureTextEx(app_font, help_line3, FONT_SIZE_SUB - 4, 1);
        DrawTextEx(app_font, help_line3,
                   (Vector2){center.x - h3_size.x / 2, help_y + 60},
                   FONT_SIZE_SUB - 4, 1, COLOR_TEXT);
    }
    else
    {
        const char *help_hint = "Press F1 for Help";
        Vector2 hint_size = MeasureTextEx(app_font, help_hint, FONT_SIZE_SUB - 4, 1);
        DrawTextEx(app_font, help_hint,
                   (Vector2){screen_w - PADDING - hint_size.x - 5, screen_h - PADDING - 20},
                   FONT_SIZE_SUB - 4, 1, COLOR_TEXT);
    }

    EndDrawing();
}

// ================================================================================
// System Safety Checks
// ================================================================================
static bool system_safety_check()
{
    if (IsWindowReady())
        return true;
    fprintf(stderr, "Error: Failed to initialize display (check X11/Wayland)\n");
    return false;
}

// ================================================================================
// Main Entry Point
// ================================================================================
int main(int argc, char *argv[])
{

#if ENABLE_DEBUG_LOG
#else
    SetTraceLogLevel(LOG_NONE);
#endif
    // Handle command line arguments
    for (int i = 1; i < argc; i++)
    {
        if (strcmp(argv[i], "-h") == 0)
        {
            printf("rlock - System Screen Locker\n");
            printf("Usage: %s [OPTIONS]\n", argv[0]);
            printf("  -h    Show this help message\n");
            return 0;
        }
#if ENABLE_DEBUG_LOG
        else if (strcmp(argv[i], "-d") == 0)
        {
            printf("Debug mode enabled (see stderr for output)\n");
        }
#endif
    }

    // Initialize Raylib
    InitWindow(0, 0, WINDOW_TITLE);
    ToggleFullscreen();
    SetTargetFPS(60);
    HideCursor();
    DisableCursor();
    SetExitKey(0); // Disable default exit key

    // Load system font (fallback to default if failed)
    default_font = GetFontDefault();
    app_font = LoadFont(SYSTEM_FONT_PATH);
    if (app_font.baseSize == 0)
    {
        fprintf(stderr, "Warning: Could not load system font - using default\n");
        app_font = default_font;
    }
    else
    {
        is_custom_font = true;
    }

    // Verify system environment
    if (!system_safety_check())
    {
        CloseWindow();
        return 1;
    }

    // Lock password memory (prevent swapping to disk)
#ifdef __linux__
    if (mlock(password, sizeof(password)) != 0)
    {
        fprintf(stderr, "Warning: Could not lock password memory (run with sufficient privileges)\n");
    }
#endif

    // Main loop
    while (is_locked)
    {
        handle_input();
        draw_ui();
    }

    // Cleanup
#ifdef __linux__
    munlock(password, sizeof(password));
#endif
    clear_password();
    EnableCursor();
    ShowCursor();
    if (is_custom_font)
        UnloadFont(app_font);
    CloseWindow();

#if ENABLE_DEBUG_LOG
    printf("System unlocked successfully\n");
#endif
    return 0;
}
